from __future__ import with_statement
try:
    import cPickle as pickle
except ImportError:
    import pickle

import sys

class LoadError(Exception):
    """Default exception that is raised when a load cannot open a pfile for reading"""

class Persistable(object):
    """
    Will add partial persistance to a class.

    Users of the implementor will be able to get and set from the dictionary
    by obj.attr and obj.attr = value.

    A save() and load() method are provided to dump/load the persistant dictionary.
    """
    def __init__(self, pfile=None, *args, **kwargs):
        super(Persistable, self).__init__(*args, **kwargs)
        self.persistant = {}
        self.pfile = pfile

    def __getattr__(self, attr):
        if attr == 'persistant':
            self.persistant = {}
            return self.persistant
        x = self.persistant.get(attr)
        if not x:
            raise AttributeError, attr
        else:
            return x

    def __setattr__(self, name, value):
        if name != 'persistant' and name in self.persistant.keys():
            self.__dict__['persistant'][name] = value
        else:
            super(Persistable, self).__setattr__(name,value)

    def __getstate__(self):
        return {'persistant':self.persistant}

    def save(self):
        if self.pfile:
            self.__save_to_file()
        else:
            return self.persistant

    def __save_to_file(self):
        with file(self.pfile, 'w') as pfile:
            pickle.dump(self.persistant, pfile)

    def load(self, exception=LoadError):
        if self.pfile:
            self.__load_from_file(exception)
        # unfortunately you have to do this manually

    def __load_from_file(self, exception):
        try:
            with file(self.pfile, 'r') as pfile:
                self.persistant.update(pickle.load(pfile))
        except IOError:
            raise exception

